import json
import platform
import argparse
from scapy.all import *
from scapy.packet import Packet
from scapy.fields import BitField, Field
from scapy.sendrecv import sendp, srp1
from scapy.utils import wrpcap

#===================== JSON解析模块 =====================
def parse_header_config(json_file):
    """解析JSON文件生成Scapy协议类"""
    with open(json_file, 'r') as f:
        config = json.load(f)
    
    field_defs = []
    for field in config:
        # 处理不同位宽的字段
        field_defs.append(
            BitField(
                name=field['name'],
                default=field['default'],
                size=field['bits']
            )
        )
    
    # 动态创建协议类
    return type('CustomProtocol', (Packet,), {'fields_desc': field_defs})

#===================== 数据包发送模块 =====================
def get_network_interfaces():
    """获取可用网络接口列表（适配Windows系统）"""
    if platform.system() == 'Windows':
        # Windows系统获取带描述的接口列表
        from scapy.arch.windows import get_windows_if_list
        return [f"{iface['description']}" 
                for iface in get_windows_if_list()]
    else:
        # Linux/Mac系统直接获取接口名称
        return get_if_list()

def send_custom_packet(packet, iface, dstAddr="127.0.0.1",timeout=2):
    """发送并接收响应数据包"""
    # 添加二层协议头确保可发送
    full_packet = Ether()/IP(dst=dstAddr)/packet
    response = srp1(full_packet, iface=iface, timeout=timeout, verbose=0)
    return full_packet, response

#===================== 数据处理模块 =====================
def save_packets(packet_list, filename,append_mode="w"):
    """保存数据包到pcap文件"""
    wrpcap(filename, packet_list,append=append_mode)

#===================== 主程序流程 =====================

def interact_with_user():
    # 解析JSON配置
    json_file = input("请输入协议定义文件路径：")
    protocol_cls = parse_header_config(json_file)

    # 用户初始化报头
    print("\n协议字段初始化：")
    pkt = protocol_cls()
    for field in pkt.fields_desc:
        current_val = getattr(pkt, field.name)
        user_input = input(f"请输入 [{field.name}] 的值（默认：{current_val}）：")
        if user_input.strip():
            setattr(pkt, field.name, int(user_input))

    # 选择网络接口
    ifaces = get_network_interfaces()
    print("\n可用网络接口：")
    for idx, iface in enumerate(ifaces):
        print(f"{idx+1}. {iface}")
    iface_idx = int(input("请选择接口编号：")) - 1

    print("\n选择目的IP地址：")
    dstAddr = input("请输入目的IP地址：(default is 127.0.0.1)")

    # 发送数据包
    sent_pkt, response = send_custom_packet(pkt, ifaces[iface_idx],dstAddr=dstAddr)

    # 保存结果
    pcap_file = input("\n请输入保存文件名（例：trace.pcap）：")
    packets = [sent_pkt]
    if response:
        packets.append(response)
    save_packets(packets, pcap_file)
    print(f"\n操作完成！已保存 {len(packets)} 个数据包到 {pcap_file}")

def shell():
        # 解析参数
    argparser = argparse.ArgumentParser(description="Custom packet sender")
    argparser.add_argument("-i", "--iface", help="Network interface name",default="eth0")
    argparser.add_argument("-j", "--json", help="JSON config file",default="protocol.json")
    argparser.add_argument("-o", "--output", help="Output pcap file",default="trace.pcap")
    argparser.add_argument("-m", "--mode",help="mode of writing pcap file",default="w")
    argparser.add_argument("-d", "--dstAddr", help="Destination IP address",default="127.0.0.1")
    argparser.add_argument("-t", "--timeout", help="Timeout for packet response",type=int,default=2,required=False)
    args = argparser.parse_args()

    proto_cls = parse_header_config(args.json)
    pkt = proto_cls()
    for field in pkt.fields_desc:
        current_val = getattr(pkt, field.name)
        user_input = input(f"请输入 [{field.name}] 的值（默认：{current_val}）：")
        if user_input.strip():
            setattr(pkt, field.name, int(user_input))
    sent_pkt,response=send_custom_packet(pkt, args.iface, dstAddr=args.dstAddr,timeout=args.timeout)
    packets = [sent_pkt]
    if response:
        packets.append(response)
    # 保存文件内容
    save_packets(packets, args.output,append_mode=args.mode)

if __name__ == "__main__":
    shell()